import os
import pandas as pd
import joblib
import json
from prophet import Prophet
from django.conf import settings


# Функция для загрузки моделей
def load_models(config_path):
    with open(os.path.join(settings.BASE_DIR, 'forecast', config_path), 'r') as f:
        config = json.load(f)
    model_paths = config['model_paths']
    models = {}
    for column, path in model_paths.items():
        models[column] = joblib.load(os.path.join(settings.BASE_DIR, 'forecast', path))
    return models, config['forecast_periods']


# Функция для прогнозирования на заданное количество недель вперед
def forecast_next_weeks(models, data, periods):
    forecasts = {}
    for column, model in models.items():
        df_prophet = data.reset_index().rename(columns={'index': 'ds', column: 'y'})
        future = model.make_future_dataframe(periods=periods, freq='W')
        forecast = model.predict(future)
        forecasts[column] = forecast[['ds', 'yhat']].tail(periods)
    return forecasts


# Функция для формирования результата в требуемом формате
def format_forecast_result(forecasts):
    result = pd.DataFrame(columns=['Год', 'Номер недели'] + list(forecasts.keys()))
    rows = []
    for column, forecast in forecasts.items():
        for i, row in forecast.iterrows():
            year, week = row['ds'].isocalendar()[:2]
            rows.append({
                'Год': year,
                'Номер недели': week,
                column: int(row['yhat'])
            })
    result = pd.DataFrame(rows)
    result = result.pivot_table(index=['Год', 'Номер недели'], aggfunc='first').reset_index()
    return result
